import time
import traceback
import signal
import sys
import socket
import threading
from server import Server
from queue import Queue
from utils import *
from common import *

logging = logger

# ==================== 主服务端 ====================
class MasterServer:
    def __init__(self, host='0.0.0.0', port=9000):
        self.host = host
        self.port = port
        self.listen_socket = None
        self.slaves: dict[str, tuple[Server, socket.socket, tuple[str,str], float]] = {}  # {client_id: (serve, socket, addr, last_active)}
        self.slaves_info: dict[str, HostInfo] = {}  # {client_id: {on_task: bool}}
        self.task_info: dict[str, list[TaskInfo]] = {}
        self.task_queue = Queue()
        # ================锁================
        self.slaves_lock = threading.RLock()
        self.slaves_info_lock = threading.RLock()
        self.task_info_lock = threading.RLock()
        # ===============其他================
        self.running = threading.Event()
        self.config = get_config("./config.cfg")
        
        # 日志配置
        # logging.basicConfig(
        #     filename='master.log',
        #     format='%(asctime)s [%(levelname)s] %(message)s',
        #     level=logging.INFO
        # )
        
    def start(self):
        """启动主服务"""
        self.running.set()
        accept_connections_thread = threading.Thread(target=self._accept_connections, daemon=True)
        accept_connections_thread.start()
        check_heartbeats_thread = threading.Thread(target=self._check_heartbeats, daemon=True)
        check_heartbeats_thread.start()
        # client_control_command_thread = threading.Thread(target=self._handle_client_control_command, daemon=True)
        # client_control_command_thread.start()
        logging.info("Master server started")
        check_heartbeats_thread.join()
        accept_connections_thread.join()
        # client_control_command_thread.join()
        logging.info("Master server stopped")

    def _accept_connections(self):
        """接受客户端连接"""
        with socket.socket() as s:
            linger = struct.pack('ii', 1, 5)  # 启用 SO_LINGER，延迟时间为 10 秒
            s.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, linger)
            s.bind((self.host, self.port))
            s.listen(5)
            self.listen_socket = s
            
            while self.running.is_set():
                try:
                    conn, addr = s.accept()
                    if is_ip_compliant(addr[0], self.config['allowable_ip']):
                        serve = Server(conn, addr[0], addr[1])
                        serve.set_processor({
                            self._process_task: [MessageType.GetTask, MessageType.TaskResult, MessageType.Exit, MessageType.Allow, MessageType.Refuse],
                            self._handle_heartbeat: [MessageType.Heartbeat,],
                            self._handle_client_control_command: [MessageType.ControlCommand,], # TODO 只有处理客户端逻辑的时候才有用
                            self._handle_register: [MessageType.Register,],
                        })
                        threading.Thread(
                            target=serve.start,
                        ).start()
                    else:
                        MessageProtocol.send(conn, {
                            'type': MessageType.Refuse,
                            'message': 'This address is not allowed to link'})
                        conn.close()
                        logging.info(f"This address [{addr}] is not allowed to link")
                except Exception as e:
                    logging.info(f"Server stop ... ...: {e}")
                    self.stop()
                    break
                
    def _update_slave_info(self, client_id, serve: Server, data):
        with self.slaves_lock:
            latest_time = time.time()
            self.slaves[client_id] = (serve, serve.socket, serve.addr, latest_time)
        self._update_slaves_info(client_id, serve.addr, latest_time, data)
    
    def _sync_task_info(self, data):
        host_info = HostInfo.deserialize(data)
        task_completed = []
        is_putong_putong = False

        with self.task_info_lock:
            for refresh_info in host_info.tasks_info:
                if refresh_info.status == TaskStatus.Finished:
                    task_completed.append(refresh_info)
                    is_putong_putong = True
                for task_info in self.task_info[refresh_info.id]:
                    if task_info.dataset_name == refresh_info.dataset_name:
                        task_info.status = refresh_info.status
                        task_info.PID = refresh_info.PID
                        task_info.details = refresh_info.details
                        break
        return is_putong_putong, task_completed

    # 监听心跳
    def _handle_heartbeat(self, serve: Server):
        while serve.running.is_set():
            msg = serve.get_message()
            client_id = msg['client_id']
            data = msg['data']
            self._update_slave_info(client_id, serve, data)
            # 根据slave中的任务信息，更新自身统计的所有任务
            is_putong_putong, task_completed = self._sync_task_info(data)
            if is_putong_putong:  # 如果slave服务器有完成的任务，需要告诉它清理，并接收新的任务
                serve.send_message(build_putong_putong([t.serialize() for t in task_completed]))
                
    # 处理客户端命令
    def _handle_client_control_command(self, serve: Server):
        while serve.running.is_set():
            try:
                msg = serve.get_message()

                command = msg.get('command')
                logger.debug(f"_handle_client_control_command {msg}")
                serve.name = 'master-client'
                serve.remote_name = 'client'
                client_id = serve.remote_name
                addr = serve.addr
                if command == 'add_tasks':
                    tasks = msg.get('tasks', [])
                    tasks = [Task.model_validate(t) for t in tasks]
                    self.dispatch_tasks(tasks)
                elif command == 'stop_server':
                    logger.info(f"[{client_id}: {addr}] Control client requested to stop the server")
                    self.stop()
                elif command == 'get_slave_info':
                    with self.slaves_info_lock:
                        slave_info = []
                        for id, host_info in self.slaves_info.items():
                            slave_info.append(host_info.serialize())
                    serve.send_message(build_slave_info(slave_info))
                    logger.info(f"Sends slave host information to the client [{client_id}: {addr}] msg: [{slave_info}")
                elif command == 'get_task_info':
                    with self.task_info_lock:
                        task_info = self.serialize_task_info()
                        MessageProtocol.send(serve.socket, {'type': MessageType.TaskInfo, 'data': task_info})
                        logger.info(f"Sends slave host information to the client [{client_id}: {addr}]")
                elif command == 'get_task_results':
                    with self.task_info_lock:
                        task_id = msg.get('id')
                        if task_id in self.task_info:
                            # TODO 先验证是否都已经完成
                            is_completed = True
                            with self.task_info_lock:
                                for task_info in self.task_info[task_id]:
                                    if task_info.status != TaskStatus.Finished \
                                    or task_info.status != TaskStatus.Error:
                                        is_completed = False
                            if is_completed:
                                data = self.get_task_results(task_id)
                                serve.send_message(build_task_result(task_id, data))
                            else:
                                serve.send_message(build_refuse(f"任务ID [{task_id}] 还未完成"))
                        else:
                            serve.send_message(build_exception(f"任务ID [{task_id}] 不存在"))
                # 可以根据需要扩展更多命令
            except Exception as e:
                MessageProtocol.send(serve.socket, build_exception(f"[{serve.name}: {serve.addr}] {e}"))
                logger.error(f"{type(e)}: {e}")
    # 处理注册信息
    def _handle_register(self, serve: Server):
        addr = (serve.addr, serve.port)
        conn = serve.socket
        
        msg = serve.get_message()
        client_id = msg['client_id']
        is_rgt = False  # 是否注册成功
        logger.info(f"start register [{client_id}: {addr}]")
        serve.name = 'master-' + client_id
        serve.remote_name = client_id
        logger.debug(f"serve remote name: {serve.remote_name}")
        with self.slaves_lock:
            if client_id not in self.slaves:
                self.slaves[client_id] = (serve, conn, addr, time.time())
                is_rgt = True
        if is_rgt:  # 注册成功
            with self.slaves_info_lock:
                self.slaves_info[client_id] = HostInfo(client_id)
            logging.info(f"Slave host [{client_id}: {addr}] registered")
            MessageProtocol.send(conn, build_allow())
            # serve.registered.set()
            logger.info(f"[{serve.remote_name}: ({serve.addr}, {serve.port})] registered successfully: {serve.registered.is_set()}")
            if not self.task_queue.empty():  # 任务不为空需要通知领取任务
                MessageProtocol.send(conn, build_new_tasks())
        else:  # 注册失败
            MessageProtocol.send(conn, build_refuse('This address is already registered'))
    # 主消息处理
    def _process_task(self, serve: Server):
        # 消息路由
        conn = serve.socket
        addr = serve.addr
        while self.running.is_set():
            msg = serve.get_message()
            msg_type = msg.get('type')
            if msg_type == MessageType.GetTask:
                logger.info(f"[{serve.name}] Rcve GetTask message")
                # 从队列中获取任务
                task = self._get_next_task()
                logger.info(f"[{serve.name}] send Task message to [{serve.remote_name}]: {task}")
                # 发送任务
                MessageProtocol.send(conn, build_task(task))
                # 更新主机状态
                msg = serve.get_message()
                with self.slaves_info_lock:
                    if msg.get('type') == MessageType.Allow:
                        self.update_tasks_info_status(task, TaskStatus.Running)
                    elif msg.get('type') == MessageType.Refuse:
                        self.update_tasks_info_status(task, TaskStatus.Error)
                    else:
                        logger.warning(f"The type of message indicating whether the task was executed successfully is incorrect: type: {msg.get('type')}")
                
            elif msg_type == MessageType.TaskResult:
                self._handle_task_result(msg)
                
            elif msg_type == MessageType.SlaveInfo:
                self.update_slave_info()
                
            elif msg_type == MessageType.Exit:
                # 检查是否还有未完成任务
                # 安全退出
                MessageProtocol.send(conn, {'type': MessageType.Allow})
                logging.info(f"[{addr}] Host exit safely")
                break
                # 拒绝退出

    def _get_next_task(self) -> Task | None:
        """获取任务（生产者-消费者核心）"""
        if not self.task_queue.empty():
            return self.task_queue.get()
        return None

    def _handle_task_result(self, result: dict):
        """处理任务结果"""
        if not result['status']:
            self.task_queue.put(result['task'])  # TODO 判断是否存在，失败重试
            logging.warning(f"Task {result['task_id']} failed")
        else:
            logging.info(f"Task {result['task_id']} success")

    def _check_heartbeats(self):
        """心跳检测"""
        logging.info("Enable heartbeat detection")
        while self.running.is_set():
            try:
                time.sleep(HEARTBEAT_CHECK_PER)
                now = time.time()
                with self.slaves_lock:
                    to_remove = [
                        cid for cid, (_, _, _, last_heartbeat) in self.slaves.items()
                        if now - last_heartbeat > HEARTBEAT_TIMEOUT
                    ]
                    for cid in to_remove:
                        # TODO 检查是否还有未完成的任务
                        serve, _, _, _ = self.slaves[cid]
                        serve.stop()
                        del self.slaves[cid]
                        with self.slaves_info_lock:
                            del self.slaves_info[cid]
                        logging.warning(f"Client {cid} heartbeat detection timeout")
            except Exception as e:
                logging.error(f"Heartbeat detection error: {e}")
        logging.info(f"Heartbeat detection stopped")

    def _add_task(self, task: Task):
        """添加任务（线程安全）"""
        self.task_queue.put(task)

    def stop(self):
        if self.running.is_set():
            """停止服务"""
            self.running.clear()
            time.sleep(5)
            self.listen_socket.close()

    def dispatch_tasks(self, tasks: list[Task]):
        """添加任务列表并给所有从服务器发送消息通知"""
        # 任务的格式{"id": task_id, "dataset": {dataset_name: [subset_name ... ...]}}
        logger.debug(f"task list inner: {tasks}")
        for task in tasks:
            # 1. 添加任务到任务队列
            self._add_task(task)
            # 2. 添加任务状态信息
            task_id = task.id
            dataset: dict = task.dataset
            model_name = task.model
            with self.task_info_lock:
                if task_id not in self.task_info.keys():
                    self.task_info[task_id] = []  # 2.1 初始化
                # 2.2 并初始化状态为排队等待
                for dataset_name, subset_name_list in dataset.items():
                    self.task_info[task_id].append(
                        TaskInfo(
                            id=task_id,
                            PID=-1,
                            dataset_name=dataset_name,
                            model=model_name,
                            subset_list=subset_name_list,
                            status=TaskStatus.Waiting,
                        )
                    )

        logging.debug(f"Task list added: {tasks}")
        send_info = []
        # 3. 通知工作节点有新任务需求
        with self.slaves_lock:
            for name, (_, conn, addr, _) in self.slaves.items():
                MessageProtocol.send(conn, build_new_tasks())
                send_info.append({
                    'name': name,
                    'addr': addr
                })
        logging.info(f"Sends a message of a new task to the servers: {send_info}")

    # =============================任务信息操作==============================
    def update_tasks_info_status(self, task: Task, status: TaskStatus):
        """根据任务设置任务相关信息的状态"""
        with self.task_info_lock:
            for name in task.dataset.keys():
                for taskinfo in self.task_info[task.id]:
                    if taskinfo.dataset_name == name:
                        taskinfo.status = status
                        break
            logger.debug(f"Update the task [{task}] status to [{status.name}], ALL: [{self.task_info[task.id]}]")

    def serialize_task_info(self, task_id=None) -> dict[str, list]:
        """获取任务状态, 不指定任务ID即获取全部的"""
        task_info = {}
        if task_id is None:
            with self.task_info_lock:
                for id, task_list in self.task_info.items():
                    new_list = []
                    for task_item in task_list:
                        new_list.append(task_item.serialize())
                    task_info[id] = new_list
        else:
            with self.task_info_lock:
                if task_id not in self.task_info.keys():
                    raise Exception(f"The task ID [{task_id}] does not exist.")
                task_list = self.task_info[task_id]
                new_list = []
                for task_item in task_list:
                    new_list.append(task_item.serialize())
                task_info[task_id] = new_list
        return task_info

    def update_tasks_details(self, task: Task):
        """更新任务状态信息的详情"""
        # 对指定的id和dataset查询详情
        # 更新任务详情
        # 更新任务状态（正在运行还是运行失败）
        pass
    # =============================slave信息操作==============================
    def _update_slaves_info(self, serve_name, addr, latest_time, data):
        try:
            # 判断消息是否可用
            if data:
                # 更新自身存储的信息
                with self.slaves_info_lock:
                    self.slaves_info[serve_name] = HostInfo.deserialize(data)
                    self.slaves_info[serve_name].addr = addr[0]
                    self.slaves_info[serve_name].port = addr[1]
                    self.slaves_info[serve_name].times = latest_time
                    logger.debug(f"update slave_info: {self.slaves_info[serve_name]}")
            else:
                logger.error(f"slave info :{data} not applicable")
        except Exception as e:
            logger.error(f"slave_info thread occur exception: {e}")
            
    # =============================客户端信息操作==============================
    def get_task_results(self, task_id: str):
        """此函数不加锁并保证接收到的id存在"""
        return {
                "results_summary": [],
                "pie_chart": [''],
                "statistical_table": [],
            }
        


def signal_handler(sig, frame):
    print("Interrupt signal received. Stopping servers...")
    if master: master.stop()
    sys.exit(0)
    
def server_start():
    # 启动主服务端
    global master
    master = MasterServer()
    master.start()

# ==================== 使用示例 ====================
if __name__ == '__main__':
    # 注册信号处理函数
    signal.signal(signal.SIGINT, signal_handler)
    server_start()